from unstructured.partition.html import partition_html
from concurrent.futures import ThreadPoolExecutor
from langchain_community.document_loaders import RecursiveUrlLoader, WebBaseLoader
from langchain_core.documents import Document
from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Annotated
import time
import os
import asyncio
import requests
import html2text
import json
import urllib.parse
from bs4 import BeautifulSoup
from langchain.tools import BaseTool, StructuredTool, tool
from datetime import datetime
import pytz
import brotli
from typing import NamedTuple
from bs4.element import Tag, NavigableString


default_header_template = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36 Edg/129.0.0.0",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*"
    ";q=0.8",
    "Accept-Language": "zh-CN,zh;q=0.5",
    "Accept-Encoding": "gzip, deflate, br",
    "Accept-Charset": "UTF-8",
    "Referer": "https://www.baidu.com/?toWww=0&redig=31327B95FBB74239837BCB5B9EBCE1D2",
    "DNT": "1",
    "cache-control": "no-cache",
    "ect": "4g",
    "cookie": "MUID=1D877AA6885A65DE11ED694E89FE64DE; MUIDB=1D877AA6885A65DE11ED694E89FE64DE; SRCHD=AF=NOFORM; SRCHUID=V=2&GUID=53C1D48576B14BF5A551244A2793F780&dmnchg=1; WLS=C=31bf5af3d2129349&N=%e6%b5%a9; ANON=A=FE69E647D1958A06AA999F8FFFFFFFFF; CortanaAppUID=8D12507101CC3F68D22D4FB5E0F6E481; MUIDV=NU=1; MMCASM=ID=3735308D6E7341A6B880D2BD914553CA; msau=id=FE69E647D1958A06AA999F8FFFFFFFFF&msa=1&aad=0; _EDGE_S=mkt=zh-cn&ui=zh-cn&SID=0FBEC179AB8E6AFF11FCD291AACD6BB6; ABDEF=V=13&ABDV=13&MRNB=1722622529964&MRB=0; ENSEARCH=BENVER=0; _UR=QS=0&TQS=0&Pn=0; SNRHOP=I=&TS=; _U=14JgolbdoI10s0A_e7Xz09aKb-va77qXdopdKE9EDa53mLbKpeU8qldh9MvlCHTdJWs_pfKb5pO62ZT916eIKo3u3z7en08JER0kfMBOyY3Xm3FOc_jk4NidekQaqxHwYRslC3mRLW8gYB89JVB3mVP7QSdgbLUURlsEODj4lR-GgucJNuvj_Vfm54pHo4f5eKM3RfVe3r696bcozhiK3Dg; ak_bmsc=0380972D59027FF80658B0514A49D002~000000000000000000000000000000~YAAQzGjBF0bXeYGSAQAAY4fVlRmWW2qzOqTVCOUUlmfn/ShJl9i6qnrOwTuMrVtnYOR6CSL4YmSGb4SYMTc95hnLtczZPQ6NInawYmROu6YENisYmNdnWI8PSwfK7KYOXuOD2EDxDHddbdAQTWkk/x27/uOUeut25uXGa+bDeZ0/qvzUIJZLl6zcXNMmTj5CO2ojJRwY6sGatk0hKDBX5BPOlxAAlEuxhVtynWfHQzIUIawHUspteougzWDogqQLP06leGeKguE+YljJOBjt7C2u8Ll/MKjYOfcn2+7Wzmrdaeq/g43f182MPDMTN6/OKn78d+NlT2p7yyR9DgX4rM5b1luR6I8bWKFBVUF/msi8jIiC16D5GwQWFkDJcvFtojkykgOVDnhA; _HPVN=CS=eyJQbiI6eyJDbiI6NCwiU3QiOjAsIlFzIjowLCJQcm9kIjoiUCJ9LCJTYyI6eyJDbiI6NCwiU3QiOjAsIlFzIjowLCJQcm9kIjoiSCJ9LCJReiI6eyJDbiI6NCwiU3QiOjAsIlFzIjowLCJQcm9kIjoiVCJ9LCJBcCI6dHJ1ZSwiTXV0ZSI6dHJ1ZSwiTGFkIjoiMjAyNC0xMC0xNlQwMDowMDowMFoiLCJJb3RkIjowLCJHd2IiOjAsIlRucyI6MCwiRGZ0IjpudWxsLCJNdnMiOjAsIkZsdCI6MCwiSW1wIjoxNCwiVG9ibiI6MH0=; SRCHUSR=DOB=20231215&T=1729090651000; MicrosoftApplicationsTelemetryDeviceId=57a0a93f-8195-4530-8ca8-52505ab502db; ai_session=tirVVszuVdXXt2AbDRWMTG|1729090668675|1729090668675; ipv6=hit=1729094270236&t=4; _Rwho=u=d&ts=2024-10-16; bm_sv=FE86F2B04282459F3556BA29CC67A8FA~YAAQzGjBF2HXeYGSAQAAE9rVlRlJgjQ7y4DdLn20mlw/CII5/G+992uib78ffjeQGWS9e3CyRe2COfP7kycMfPds3DP4gGdnp/o8aeo1G5WtxTIcLg4U0mRzcxtgvoRs6j5We6eD3Y6bDPI9O8oAubN4ixohMeBXrBJsreyD98EmvOPXD5xU1ozZdRDe10H0cpyyKrqwE/35TscPwc2743fF7v7Ti8VrcEw1y1k+tYft2ooXwhzqJwbtHEy8qg==~1; USRLOC=HS=1&ELOC=LAT=22.631732940673828|LON=114.06842041015625|N=%E9%BE%99%E5%B2%97%E5%8C%BA%EF%BC%8C%E5%B9%BF%E4%B8%9C%E7%9C%81|ELT=9|&BID=MjQxMDE2MjIyNzU5X2QwZTZkNDQxMjY0M2M1ZGMwOWJmMTkyMGI2YjRmZTVhMjM0OTNlMGYyMmFlOWNhZWYwODQzNDQwY2I2NDJiYjg=; _C_ETH=1; GC=n8GR8U4BI9A6FsJ3CVQzzuuoC1XpZIwRxl0F_wa2QpMEjGhc4moKx1IidYeJgWQd1gWDH0RGvATY1QoS5UC1ZQ; _RwBf=mta=0&rc=316&rb=316&gb=0&rg=0&pc=311&mtu=0&rbb=0.0&g=0&cid=&clo=0&v=2&l=2024-10-16T07:00:00.0000000Z&lft=0001-01-01T00:00:00.0000000&aof=0&ard=0001-01-01T00:00:00.0000000&o=0&p=BINGCOPILOTWAITLIST&c=MR000T&t=9663&s=2023-03-29T02:15:21.0070277+00:00&ts=2024-10-16T14:57:55.8600070+00:00&rwred=0&wls=2&wlb=0&wle=0&ccp=2&lka=0&lkt=0&aad=0&TH=&e=Zom_wXxYpA_1vSD1EfwxQk9ioUyCT8FNIQSRjZMqERxEffeTC04nGSbU3mtQsKJBJ-OSE1zo-Kwcoobm36nlXw&A=&rwdbt=-62135539200&rwflt=1726622438&cpt=0&rwaul2=0; _SS=SID=0FBEC179AB8E6AFF11FCD291AACD6BB6&R=316&RB=316&GB=0&RG=0&RP=311&PC=U531&OCID=winp2fptaskbar; SRCHHPGUSR=SRCHLANG=zh-Hans&PV=15.0.0&BZA=0&BRW=S&BRH=M&CW=1047&CH=925&SCW=1164&SCH=2088&DPR=2.0&UTC=480&DM=0&EXLTT=31&HV=1729090675&PRVCW=1920&PRVCH=925&PR=2&IG=9953F218264A465A8D36E466894ED470&THEME=0&WEBTHEME=0;Hm_lvt_2bc5b38b6ac7f27efef8a3ad89d1a6cc=1729155438; Hm_lpvt_2bc5b38b6ac7f27efef8a3ad89d1a6cc=1729155438; HMACCOUNT=5D045E531F54DEA4; uniqueId=6350c96d6a3485863856358725e66edf",
    "sec-ch-ua": '"Microsoft Edge";v="129", "Not=A?Brand";v="8", "Chromium";v="129"',
    "x-client-data": "eyIxIjoiMCIsIjEwIjoiXCJwQ3psbmFmcU8rWWlUY2RaM3RoY2U1dUczakt2b3RYM0xtUG5SMHBSQzUwPVwiIiwiMiI6IjAiLCIzIjoiMCIsIjQiOiI0Nzc2OTgyMDk4NjQxNTI4MDQyIiwiNSI6IlwiRUJlSW9uRDFtbytBNUxHTzVtVmdqazBseWhENEZpNW03ZmdYcm03T3U3Yz1cIiIsIjYiOiJzdGFibGUiLCI3IjoiMzQzNTk3MzgzNjgwNCIsIjkiOiJkZXNrdG9wIn0=",
    "sec-ch-ua-platform": "Windows",
    "sec-fetch-dest": "document",
    "sec-ch-ua-mobile": "?0",
    "sec-fetch-mode": "navigate",
    "sec-fetch-site": "same-origin",
    "sec-fetch-user": "?1",
    "Connection": "keep-alive",
    "Upgrade-Insecure-Requests": "1",
}


session = requests.Session()
session.verify = True
session.headers = dict(default_header_template)


def get_bing_search_result_json(query: str, num_results: int = 4, is_gitee: bool = False) -> dict:
    """获取必应搜索结果,字典格式"""

    try:
        res = session.get(
            f"https://cn.bing.com/search?q={urllib.parse.quote_plus((query+' site:gitee.com') if is_gitee else query)}&setlang=zh-Hans&ghacc=0&ghsh=0&mkt=zh-CN", headers=default_header_template, timeout=10)
        res.raise_for_status()
        res.encoding = 'utf-8'
        results = {'organic_results': list(
            dict()), 'status': res.ok}
        soup = BeautifulSoup(res.text, "html.parser")
        print("抓取成功")
        b_results = soup.find('ol', id="b_results")
        if b_results:
            if b_results and isinstance(b_results, Tag):
                for pos, elm in enumerate(b_results.find_all('li', class_='b_algo')):
                    source = title = link = summary = None
                    if elm.find('div', class_="tptt"):
                        source = elm.find('div', class_="tptt").text
                    if elm.find('h2'):
                        title = elm.find('h2').text
                        if elm.find('h2').find('a'):
                            link = elm.find('h2').find('a')['href']
                    if elm.find('p'):
                        summary = elm.find('p').text
                    results['organic_results'].append(
                        {'position': pos+1, 'source': source, 'title': title, 'link': link, 'summary': summary})
                    if pos+1 == num_results:
                        break
            results['number_of_results'] = len(results['organic_results'])
        return results
    except requests.exceptions.HTTPError as err:
        print(f"请求失败，错误信息: {err}")
        return {'organic_results': [], 'status': False}


def fetch_page_content(url: str) -> str:
    """抓取指定 url 的 页面信息。"""
    if not url:
        return "URL 不能为空"
    search_time = time.time()
    try:
        res = session.get(url, timeout=10,)
        res.encoding = 'utf-8'
        res.raise_for_status()
        pure_content = "\n".join([str(el) for el in partition_html(
            text=res.text, skip_headers_and_footers=True)])
        print(f"抓取 {url} 耗时", time.time() - search_time)
        return pure_content[:6000] or "抓取内容为空"
    except Exception as e:
        print(e)
        return '网络异常'


async def fetch_page_content_async(url, executor):
    loop = asyncio.get_running_loop()
    try:
        result = await loop.run_in_executor(executor, fetch_page_content, url)
        return result
    except Exception as e:
        print(f"抓取 {url} 时发生错误: {e}")
        return '抓取失败'


@tool
def get_now_time(p: str) -> str:
    """获取当前时间"""
    tz = pytz.timezone('Asia/Shanghai')
    now = datetime.now(tz)
    return now.strftime("%Y-%m-%d %H:%M:%S")


@tool
def get_search_result(query: Annotated[str, "搜索关键词"],
                      num_results: Annotated[int,
                                             "搜索数量，复杂度越高值应越大。范围 3-6"] = 3,
                      ) -> str:
    """联网搜索任何信息"""
    # headers=Optional(default header will be used), max_retries=Optional(default: 3), num_results=Optional(default: 10)
    try:
        bing_results = get_bing_search_result_json(
            query, num_results=num_results)
        print("必应搜索列表", bing_results)

        async def main(links):
            with ThreadPoolExecutor(max_workers=8) as executor:
                tasks = [fetch_page_content_async(
                    link, executor) for link in links]
                return await asyncio.gather(*tasks, return_exceptions=True)

        links_list = [item['link'] for item in bing_results['organic_results']]
        summary_list = [{"title": item['title'], 'summary': item['summary'], 'url': item['link']}
                        for item in bing_results['organic_results']]
        print('链接列表', links_list)
        if (len(links_list) > 0):
            search_time = time.time()
            search_results_list = asyncio.run(main(links_list))
            print("总抓取耗时:", time.time() - search_time)
            search_results_str = ""
            for r in search_results_list:
                if isinstance(r, str):
                    search_results_str += r
            # print("搜索结果:", search_results_str)
            print("抓取总字符数:", len(search_results_str))
            return "摘要:"+str(summary_list) + "\n" + search_results_str[:8000] or "联网搜索结果为空"
        return "联网搜索结果为空"
    except Exception as e:
        print("网络异常", e)
        return '网络异常'


@tool
def python_code_exec(code: Annotated[str, '安全的 python 代码、表达式']) -> str:
    """执行 python 代码、数学计算。可利用 python 编程解决的问题，你可以自由编写 python 代码，结果赋值到 result 变量。禁止任何不安全代码，包括但不限于：获取任何系统信息、环境变量、文件操作等等"""
    local_vars = {}
    try:
        exec(code, {}, local_vars)
        return f"结果是: {str(local_vars)}"
    except Exception as e:
        return f"执行失败: {str(e)}"


tools_list = [get_search_result, get_now_time, python_code_exec]
